#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>

#define HEADER_SIZE 32
#define MAX_FIRMWARE_TYPE_LENGTH 10

typedef struct {
    unsigned char magic1;
    unsigned char magic2;
    unsigned char firmware_count;
    unsigned char major_version;
    unsigned char minor_version;
    unsigned char reserved[11];
    unsigned int firmware_length;
    unsigned short firmware_checksum;
    char firmware_type[MAX_FIRMWARE_TYPE_LENGTH];
} OtaHeader;

void print_help() {
    printf("Usage: otapack.exe [-h] [-p file type version] [-u file] [-v file]\\n");
    printf("Options:\\n");
    printf("  -h\t\tPrint this help message\\n");
    printf("  -p\t\tPack image file into OTA upgrade package\\n");
    printf("  -u\t\tUnpack OTA upgrade package into bin image file\\n");
    printf("  -v\t\tView information of OTA upgrade file\\n");
}

int hex_to_bin(const char *hex_file, const char *bin_file) {
    FILE *hex_fp = fopen(hex_file, "r");
    if (!hex_fp) {
        perror("Failed to open hex file");
        return 1;
    }

    FILE *bin_fp = fopen(bin_file, "wb");
    if (!bin_fp) {
        perror("Failed to open bin file");
        fclose(hex_fp);
        return 1;
    }

    char line[256];
    while (fgets(line, sizeof(line), hex_fp)) {
        if (line[0] == ':') {
            int record_length = (line[1] << 4) | (line[2] & 0xF);
            int offset = (line[3] << 12) | (line[4] << 8) | (line[5] << 4) | (line[6] & 0xF);
            int record_type = (line[7] << 4) | (line[8] & 0xF);

            if (record_type == 0) {
                for (int i = 0; i < record_length; i++) {
                    unsigned char byte = (line[9 + i * 2] << 4) | (line[10 + i * 2] & 0xF);
                    fseek(bin_fp, offset + i, SEEK_SET);
                    fwrite(&byte, 1, 1, bin_fp);
                }
            }
        }
    }

    fclose(hex_fp);
    fclose(bin_fp);
    return 0;
}

int cly_to_bin(const char *cly_file, const char *bin_file) {
    FILE *cly_fp = fopen(cly_file, "r");
    if (!cly_fp) {
        perror("Failed to open cly file");
        return 1;
    }

    FILE *bin_fp = fopen(bin_file, "wb");
    if (!bin_fp) {
        perror("Failed to open bin file");
        fclose(cly_fp);
        return 1;
    }

    char line[256];
    fgets(line, sizeof(line), cly_fp); // Read the first line (program length)
    int program_length = (line[0] << 24) | (line[1] << 16) | (line[2] << 8) | (line[3] & 0xFF);

    while (fgets(line, sizeof(line), cly_fp)) {
        for (int i = 0; i < strlen(line) / 2; i++) {
            unsigned char byte = (line[i * 2] << 4) | (line[i * 2 + 1] & 0xF);
            fwrite(&byte, 1, 1, bin_fp);
        }
    }

    fclose(cly_fp);
    fclose(bin_fp);
    return 0;
}

int main(int argc, char *argv[]) {
    int opt;
    char *file = NULL;
    char *type = NULL;
    char *version = NULL;
    int unpack = 0;
    int view = 0;

    while ((opt = getopt(argc, argv, "hp:u:v:")) != -1) {
        switch (opt) {
            case 'h':
                print_help();
                return 0;
            case 'p':
                file = optarg;
                type = argv[optind++];
                version = argv[optind++];
                break;
            case 'u':
                file = optarg;
                unpack = 1;
                break;
            case 'v':
                file = optarg;
                view = 1;
                break;
            default:
                fprintf(stderr, "Unknown option\\n");
                print_help();
                return 1;
        }
    }

    if (file == NULL) {
        fprintf(stderr, "No file specified\\n");
        print_help();
        return 1;
    }

    if (unpack) {
        // Implement unpacking logic here
        printf("Unpacking %s\\n", file);

        FILE *ota_fp = fopen(file, "rb");
        if (!ota_fp) {
            perror("Failed to open ota file");
            return 1;
        }

        OtaHeader header;
        if (fread(&header, sizeof(OtaHeader), 1, ota_fp) != 1) {
            perror("Failed to read OTA header");
            fclose(ota_fp);
            return 1;
        }

        if (header.magic1 != 0xAA || header.magic2 != 0x55) {
            fprintf(stderr, "Invalid OTA file header\\n");
            fclose(ota_fp);
            return 1;
        }

        char bin_file[256];
        snprintf(bin_file, sizeof(bin_file), "%s.bin", file);

        FILE *bin_fp = fopen(bin_file, "wb");
        if (!bin_fp) {
            perror("Failed to open bin file");
            fclose(ota_fp);
            return 1;
        }

        unsigned char byte;
        for (unsigned int i = 0; i < header.firmware_length; i++) {
            if (fread(&byte, 1, 1, ota_fp) != 1) {
                perror("Failed to read firmware image");
                fclose(ota_fp);
                fclose(bin_fp);
                return 1;
            }
            fwrite(&byte, 1, 1, bin_fp);
        }

        fclose(ota_fp);
        fclose(bin_fp);

        printf("Unpacked %s into %s\\n", file, bin_file);
    } else if (view) {
        // Implement viewing logic here
        printf("Viewing %s\\n", file);

        FILE *ota_fp = fopen(file, "rb");
        if (!ota_fp) {
            perror("Failed to open ota file");
            return 1;
        }

        OtaHeader header;
        if (fread(&header, sizeof(OtaHeader), 1, ota_fp) != 1) {
            perror("Failed to read OTA header");
            fclose(ota_fp);
            return 1;
        }

        if (header.magic1 != 0xAA || header.magic2 != 0x55) {
            fprintf(stderr, "Invalid OTA file header\\n");
            fclose(ota_fp);
            return 1;
        }

        printf("OTA File Information:\\n");
        printf("Firmware Count: %d\\n", header.firmware_count);
        printf("Major Version: %d\\n", header.major_version);
        printf("Minor Version: %d\\n", header.minor_version);
        printf("Firmware Length: %u\\n", header.firmware_length);
        printf("Firmware Checksum: %hu\\n", header.firmware_checksum);
        printf("Firmware Type: %s\\n", header.firmware_type);

        fclose(ota_fp);
    } else {
        // Implement packing logic here
        printf("Packing %s as %s version %s\\n", file, type, version);

        char bin_file[256];
        snprintf(bin_file, sizeof(bin_file), "%s.bin", file);

        if (strstr(file, ".hex") != NULL) {
            hex_to_bin(file, bin_file);
        } else if (strstr(file, ".cly") != NULL) {
            cly_to_bin(file, bin_file);
        } else if (strstr(file, ".bin") == NULL) {
            strcpy(bin_file, file);
        }

        // Calculate firmware length
        FILE *bin_fp = fopen(bin_file, "rb");
        if (!bin_fp) {
            perror("Failed to open bin file");
            return 1;
        }

        fseek(bin_fp, 0, SEEK_END);
        unsigned int firmware_length = ftell(bin_fp);
        fseek(bin_fp, 0, SEEK_SET);

        // Calculate firmware checksum
        unsigned short firmware_checksum = 0;
        unsigned char byte;
        while (fread(&byte, 1, 1, bin_fp) == 1) {
            firmware_checksum += byte;
        }

        fclose(bin_fp);

        // Build OTA header
        OtaHeader header;
        header.magic1 = 0xAA;
        header.magic2 = 0x55;
        header.firmware_count = 1;
        header.major_version = atoi(version);
        header.minor_version = atoi(strchr(version, '.') + 1);
        memset(header.reserved, 0, sizeof(header.reserved));
        header.firmware_length = firmware_length;
        header.firmware_checksum = firmware_checksum;
        strncpy(header.firmware_type, type, MAX_FIRMWARE_TYPE_LENGTH);

        // Write OTA file
        char ota_file[256];
        snprintf(ota_file, sizeof(ota_file), "%s_ota.bin", file);

        FILE *ota_fp = fopen(ota_file, "wb");
        if (!ota_fp) {
            perror("Failed to open ota file");
            return 1;
        }

        fwrite(&header, sizeof(OtaHeader), 1, ota_fp);

        // Write firmware image
        bin_fp = fopen(bin_file, "rb");
        if (!bin_fp) {
            perror("Failed to open bin file");
            fclose(ota_fp);
            return 1;
        }

        while (fread(&byte, 1, 1, bin_fp) == 1) {
            fwrite(&byte, 1, 1, ota_fp);
        }

        fclose(bin_fp);
        fclose(ota_fp);

        printf("Packed %s into %s\\n", file, ota_file);
    }

    return 0;
}
